# 环节4：构建向量数据库（基于Chroma，兼容通义千问1536维向量）
# 导入必要库（含前序环节依赖）
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_community.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter  # 按你LangChain版本选路径
from langchain_community.vectorstores import Chroma  # 新版Chroma导入路径
import os  # 用于处理文件路径

from langchain_rag.rag.document_loading import load_document
from langchain_rag.rag.document_splitting import split_documents
from langchain_rag.rag.tongyi_text_vectorization import initialize_tongyi_embeddings
from rag_test.langchain.text_vectorization import vectorize_documents_with_tongyi


def load_split_vectorize_pdf(pdf_path, api_key):
    """
    一站式完成：PDF加载→文档分割→通义千问向量化
    返回：分割后的文档对象列表、对应的向量列表
    """
    # 1. 加载PDF
    pages = load_document(pdf_path)
    split_docs = split_documents(pages)
    embeddings = initialize_tongyi_embeddings(api_key)
    vectors, metadata_list = vectorize_documents_with_tongyi(embeddings, split_docs)
    return split_docs, vectors, metadata_list


def build_chroma_vector_db(split_docs, vectors, db_path="./tongyi_chroma_db"):
    if os.path.exists(db_path):
        print(f"⚠️  发现旧数据库：{db_path}，将自动删除并重建")
        import shutil
        shutil.rmtree(db_path)  # 递归删除文件夹

    # 3. 构建Chroma数据库（核心步骤）
    try:
        # 创建空的Chroma数据库
        db = Chroma(
            embedding_function=None,  # 关键：已手动生成向量，设为None
            persist_directory=db_path  # 持久化存储路径（本地文件夹）
        )

        # 手动添加预生成的向量（通义千问的1024维向量）
        # 注意：Chroma要求向量格式为list[list[float]]，与通义千问输出格式一致
        texts = [doc.page_content for doc in split_docs]
        metadatas = [doc.metadata for doc in split_docs]
        ids = [f"doc_{i}" for i in range(len(split_docs))]
        
        db._collection.add(
            documents=texts,
            embeddings=vectors,  # 通义千问生成的向量列表
            metadatas=metadatas,
            ids=ids
        )

        # 持久化数据库（确保数据写入本地文件，防止程序退出后丢失）
        db.persist()
        print(f"✅ Chroma向量数据库构建完成！")

        return db

    except Exception as e:
        raise Exception(f"❌ 数据库构建失败：{str(e)}（可能是权限不足，检查文件夹读写权限）")


if __name__ == "__main__":
    file_path = "平安保险用户手册.pdf"
    api_key = "sk-6434b554122c4b049ceb805d703f695b"
    split_docs, vectors, metadata_list = load_split_vectorize_pdf(file_path, api_key)
    ##print("split_docs", split_docs)
    #print("vectors", vectors)
    #print("metadata_list", metadata_list)

    build_chroma_vector_db(split_docs,vectors)
